import os
import shutil

from gxl_ai_utils.config.gxl_config import GxlNode
from gxl_ai_utils.utils import utils_file
import logging
from gxl_ai_utils.thread.my_thread import GxlDynamicThreadPool

config = GxlNode.get_config_from_yaml("./config.yaml")
input_root_dir = config.input_dir
output_root_dir = config.output_dir


def handle_in_one_gpu(gpu_id, input_dir_list, thread_num_pre_gpu=1):
    for dir in input_dir_list:
        logging.info(f'开始处理{dir}目录')
        wav_scp = utils_file.join_path(dir, "wav.scp")
        text_scp = utils_file.join_path(dir, "text")
        if not os.path.exists(wav_scp) or not os.path.exists(text_scp):
            logging.info(f'skip {dir}, the dir do not have wav.scp or text file')
            continue
        output_dir = utils_file.join_path(output_root_dir, dir)
        if os.path.exists(output_dir):
            shutil.rmtree(output_dir)
        os.makedirs(output_dir, exist_ok=True)
        wav_dic = utils_file.load_dic_from_scp(wav_scp)
        temp_list_file = utils_file.join_path(output_dir, "temp.list")
        wav_list = list(wav_dic.keys())
        utils_file.write_list_to_file(wav_list, temp_list_file)

        os.system(
            f"bash ./pys/music_denoise/run_test.sh {gpu_id} {output_dir} {temp_list_file} {thread_num_pre_gpu}")


def main():
    res_dir_list = []
    for root, dirs, files in os.walk(input_root_dir):
        for dir in dirs:
            res_dir_list.append(utils_file.join_path(root, dir))
        break
    gpu_num = 2
    gpu_ids = [1, 2]
    thread_num_pre_gpu = 10
    thread_runner = GxlDynamicThreadPool()
    for i in range(gpu_num):
        dir_list = res_dir_list[i::gpu_num]
        thread_runner.add_thread(handle_in_one_gpu, [gpu_ids[i], dir_list, thread_num_pre_gpu])
    thread_runner.start()

    print('all done，do by gxl!')
